package com.math.aistudykbr.config;

import jakarta.annotation.PreDestroy;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.concurrent.*;

@Component
@Slf4j
public class SSEConnectManager {
    private static final ConcurrentHashMap<String, SseEmitter> CONNECTION_POOL = new ConcurrentHashMap<>();
    private final ScheduledExecutorService heartbeatExecutor =
            Executors.newScheduledThreadPool(4);

    public void addConnection(String userId, SseEmitter emitter) {
        // 1. 设置超时和清理回调
        emitter.onTimeout(() -> {
            log.warn("连接超时: {}", userId);
            removeConnection(userId, emitter);
        });
        emitter.onCompletion(() -> {
            log.info("连接关闭: {}", userId);
            removeConnection(userId, emitter);
        });

        // 2. 启动心跳
        ScheduledFuture<?> heartbeatFuture = startHeartbeat(userId, emitter);

        // 3. 存储连接（原子性替换）
        CONNECTION_POOL.compute(userId, (k, oldEmitter) -> {
            if (oldEmitter != null) {
                oldEmitter.completeWithError(new IllegalStateException("被新连接替换"));
            }
            return emitter;
        });
    }

    private ScheduledFuture<?> startHeartbeat(String userId, SseEmitter emitter) {
        return heartbeatExecutor.scheduleAtFixedRate(() -> {
            try {
                emitter.send(SseEmitter.event()
                        .data("{\"type\":\"heartbeat\", \"time\": " + System.currentTimeMillis() + "}")
                );
            } catch (IOException e) {
                log.error("心跳失败，移除连接: {}", userId, e);
                removeConnection(userId, emitter);
            }
        }, 0, 5, TimeUnit.MINUTES);
    }

    public void removeConnection(String userId, SseEmitter emitter) {
        CONNECTION_POOL.remove(userId, emitter);
    }

    public SseEmitter getConnection(String userId) {
        return CONNECTION_POOL.get(userId);
    }

    @PreDestroy
    public void shutdown() {
        heartbeatExecutor.shutdown(); // 应用关闭时清理线程池
    }
}